#!/usr/bin/python3
# Prof. Shi Lei's Dynamic Programming Algorithm
# by Shi Lei & Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import numpy
import matplotlib
matplotlib.use('Agg'); # plot and save, never showing

import matplotlib.pyplot as plt

def calcXvector(clussprs):
  '''
# for a paper p:
#  x(p)=1/sqrt(|#c(p)|)
# input:
#  clussprs: cluster seperators: [0,p_1,p_2,...,p_n,|#P|]
  '''
  xvec=numpy.zeros(clussprs[-1])+1.;
  for c in range(1,len(clussprs)):
   xvec[range(clussprs[c-1],clussprs[c])]=1./numpy.sqrt(clussprs[c]-clussprs[c-1]);
  return xvec;

def monitorXandErr(eigvec,clussprs):
  '''
# monitoring x vector and error function
  '''
  xvec=calcXvector(clussprs);
  reerr=1-xvec.dot(eigvec)**2/xvec.dot(xvec)/eigvec.dot(eigvec);
  return xvec,reerr;

def visualizeDivide(eigvec,clussprs):
  '''
# visualize dividing progress
  '''
  mappedxvec=calcXvector(clussprs);
  mappedeigv=eigvec;
  maponexvec=mappedxvec/numpy.sqrt(len(clussprs)-1);
  plt.clf();
  plt.plot(range(0,len(eigvec)),mappedeigv,color="green",linewidth=2,linestyle='--',label='eigenvector');
  plt.plot(range(0,len(eigvec)),maponexvec,color="blue", linewidth=2,linestyle='-',label='x vector');
  plt.grid();
  plt.xlabel('paper index (sorted by eigvec)');
  plt.ylabel('vector component');
  plt.legend(loc='upper left');
  plt.savefig('../out/xveig'+str(len(clussprs)-1)+'.png');
  # plt.show(); # don't forget to commentize "matplotlib.use('Agg');" in the beginning

def clussprs2clusters(clussprs,sortmap):
  '''
# generate clusters set from cluster seperator data
  '''
  def _idxfunc(x):
    return x[1];
  clusters=[list(set([_idxfunc(sortmap[p]) for p in range(clussprs[c],clussprs[c+1])])) for c in range(0,len(clussprs)-1)];
  return clusters;

def calcClusters(eigvec, clusnum=20, monitor=False):
  '''
# set paper clusters to make b=sum(a|c)/sqrt(#c) approching sum(a)/sum(sqrt(#c))
# Algorithm by Shi Lei <shijim@gmail.com>
# Program by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  eigenvector
#  clusnum: max number of clusters
#  monitor: whether visualize the progress
# output:
#  snapshot (last history state detail): clusters set [[paper indices],...]
  '''
  dimension=len(eigvec);

  def _keyfunc(x):
   return x[0];
  sortmap=sorted([(eigvec[p],p) for p in range(0,len(eigvec))],key=_keyfunc); # (eigen_value,origin_order)
  #sortbck=sorted([(sortmap[p][1],p) for p in range(0,len(eigvec))],key=_keyfunc); # (origin_order,eigen_order)
  mappedeigv=numpy.float64([e[0] for e in sortmap]);
  numpyeigvec=numpy.array(mappedeigv); # different from list when slicing

  if (monitor):
   print('Dynamic Programming Algorithm:'+'\n  paper number: '+str(dimension)+'\n  set cluster number: '+str(clusnum));
  if (dimension<clusnum):
   print('Setting Cluster Number is smaller than Current Cluster Number, reset to equal.');
   clusnum=dimension;

  # Best 2D array initialization: http://stackoverflow.com/a/2398187
  # clussprs=[x[:] for x in [[None]*dimension]*clusnum]; # best seperators (considering recursive relation, only last seperator recorded)
  clusobjs=[x[:] for x in [[0]*dimension]*(clusnum+1)]; # best objectives: M(i,j)
  eigvsums=[None]*dimension; # eigenvector component sum: Q(0,i)

  # boundary conditions
  eigvtmp=0;
  for i in range(0,dimension):
   eigvtmp+=numpyeigvec[i];
   eigvsums[i]=eigvtmp;
   clusobjs[1][i]=eigvtmp/numpy.sqrt(i+1);
  
  for j in range(1,clusnum+1):
   clusobjs[j][j-1]=eigvsums[j-1];
         
  for j in range(2,clusnum):   	  
   if monitor:   
     print('Step '+str(j-1)+'/'+str(clusnum-1));
   
   eigvsumstmp=[0]*dimension;
   clusobjtmp=[0]*dimension;
   
   for i in range(j,dimension-1):
    
    max_f = 0;
    
    #pre-compute Q(i',i), F(i',i), i' in [0,n)
    if (i==j):
    	eigvsumstmp[0] = eigvsums[i];
    	for itmp in range(1,i+1):
    		eigvsumstmp[itmp] = eigvsumstmp[itmp-1] - numpyeigvec[itmp-1];
    else:
    	for itmp in range(0,i+1):
    		eigvsumstmp[itmp] = eigvsumstmp[itmp] + numpyeigvec[i];
    
    for itmp in range(0,i+1):
    	clusobjtmp[itmp] = eigvsumstmp[itmp]/numpy.sqrt(i-itmp+1);    	    		    
    
    for itmp in range(j-1, i):
    	cur_f = clusobjs[j-1][itmp] + clusobjtmp[itmp+1];
    	
    	if (cur_f > max_f):
    		max_f = cur_f;
    		clusobjs[j][i] = cur_f;
    	    
  clussprs = [0]*(clusnum+1);
  clussprs[0] = 0;
  clussprs[clusnum] = dimension;
    
  i=dimension-1;
  eigvsumstmp=[0]*dimension;
  clusobjtmp=[0]*dimension;
  
  for j in range(clusnum,1,-1):
  	max_f = 0;
  	
  	# pre-compute Q(i',i), F(i',i), i' in [0,n)
  	if (i==(dimension-1)):
  		eigvsumstmp[0] = eigvsums[i];
  		for itmp in range(1,i+1):
  			eigvsumstmp[itmp] = eigvsumstmp[itmp-1] - numpyeigvec[itmp-1];    		    
  	else:
  		# compute the sum of eigvector in (i+pi_{j+1}]
  		tmpsum = 0;
  		for itmp in range(clussprs[j],clussprs[j+1]):
  			tmpsum+=numpyeigvec[itmp];
  		for itmp in range(0,i+1):
  			eigvsumstmp[itmp] = eigvsumstmp[itmp] - tmpsum;
    		  	
  	for itmp in range(0,i+1):
  		clusobjtmp[itmp] = eigvsumstmp[itmp]/numpy.sqrt(i-itmp+1);    	
  	
  	for itmp in range(j-1, i):
  		cur_f = clusobjs[j-1][itmp] + clusobjtmp[itmp+1];
  		if (cur_f > max_f):
  			max_f = cur_f;
  			clussprs[j-1] = itmp+1;
  			
  	i = i - (clussprs[j] - clussprs[j-1]);
  	if ((i+1) != clussprs[j-1]):
  		raise Exception('ERROR: not consistent cluster separators!!');        	          	  
  
  numvec=numpy.array([(clussprs[it+1]-clussprs[it]) for it in range(0,clusnum)]);
  sqrtnumvec=numpy.sqrt(numvec);
  eigclussum=numpy.array([sum(numpyeigvec[clussprs[it]:clussprs[it+1]]) for it in range(0,clusnum)]);
  flowdotproduct=sum(eigclussum/sqrtnumvec);
  
  clusters=clussprs2clusters(clussprs,sortmap);
    
  return clusters;

